import { Injectable, OnModuleInit } from '@nestjs/common'
import * as WebSocket from 'ws'
import { ConfigService } from '@nestjs/config'
import { EventEmitter } from 'events'

@Injectable()
export class ComfyService implements OnModuleInit {
    private ws: WebSocket | null = null
    private readonly wsUrl: string
    private readonly eventEmitter = new EventEmitter()
    private tasks = new Map<string, any>()
    private reconnectAttempts = 0
    private readonly maxReconnectAttempts = 5
    private readonly reconnectInterval = 5000 // 5秒

    constructor(private configService: ConfigService) {
        this.wsUrl = this.configService.get<string>('COMFYUI_WS_URL') || 'ws://localhost:8188/ws'
    }

    async onModuleInit() {
        await this.connect()
    }

    async startGeneration(params: any) {
        const taskId = Math.random().toString(36).substring(7)
        this.tasks.set(taskId, { params, status: 'pending' })
        await this.queuePrompt(params)
        return taskId
    }

    async saveWorkflow(workflow: any) {
        // 实现工作流保存逻辑
        return { id: Math.random().toString(36).substring(7) }
    }

    async getWorkflow(id: string) {
        // 实现获取工作流逻辑
        return { id, workflow: {} }
    }

    onProgress(taskId: string, callback: (progress: any) => void) {
        this.eventEmitter.on(`progress:${taskId}`, callback)
    }

    onComplete(taskId: string, callback: (result: any) => void) {
        this.eventEmitter.on(`complete:${taskId}`, callback)
    }

    onError(taskId: string, callback: (error: any) => void) {
        this.eventEmitter.on(`error:${taskId}`, callback)
    }

    private async connect() {
        return new Promise<void>((resolve, reject) => {
            try {
                console.log(`Attempting to connect to WebSocket at ${this.wsUrl}`)
                this.ws = new WebSocket(this.wsUrl)

                this.ws.on('open', () => {
                    console.log('Connected to ComfyUI WebSocket')
                    this.reconnectAttempts = 0
                    resolve()
                })

                this.ws.on('message', (data: WebSocket.Data) => {
                    try {
                        const message = JSON.parse(data.toString())
                        this.handleMessage(message)
                    } catch (error) {
                        console.error('Failed to parse WebSocket message:', error)
                    }
                })

                this.ws.on('error', (error) => {
                    console.error('WebSocket error:', error)
                    if (this.ws?.readyState === WebSocket.CLOSED) {
                        this.handleReconnect()
                    }
                })

                this.ws.on('close', () => {
                    console.log('WebSocket connection closed')
                    this.ws = null
                    this.handleReconnect()
                })
            } catch (error) {
                console.error('Failed to connect to WebSocket:', error)
                this.handleReconnect()
                reject(error)
            }
        })
    }

    private async handleReconnect() {
        if (this.reconnectAttempts < this.maxReconnectAttempts) {
            this.reconnectAttempts++
            console.log(`Attempting to reconnect (${this.reconnectAttempts}/${this.maxReconnectAttempts})...`)
            setTimeout(() => {
                this.connect().catch(error => {
                    console.error('Reconnection failed:', error)
                })
            }, this.reconnectInterval * this.reconnectAttempts)
        } else {
            console.error('Max reconnection attempts reached')
        }
    }

    private handleMessage(message: any) {
        const taskId = message.taskId || 'default'

        switch (message.type) {
            case 'progress':
                this.eventEmitter.emit(`progress:${taskId}`, message.data)
                break
            case 'executed':
                this.eventEmitter.emit(`complete:${taskId}`, message.data)
                this.tasks.delete(taskId)
                break
            case 'error':
                this.eventEmitter.emit(`error:${taskId}`, message.data)
                this.tasks.delete(taskId)
                break
            default:
                console.log('Unknown message type:', message)
        }
    }

    async sendMessage(message: any) {
        if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
            await this.connect()
        }
        this.ws?.send(JSON.stringify(message))
    }

    async queuePrompt(workflow: any) {
        return this.sendMessage({
            type: 'queue_prompt',
            data: workflow
        })
    }

    async getStatus() {
        if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
            return { status: 'disconnected' }
        }
        return { status: 'connected' }
    }
}